# ================== 回测函数 ==================
import os

import numpy as np
import pandas as pd
import torch
import backtrader  as bt
from matplotlib import pyplot as plt

from 多因子.框架 import 模型训练
from 多因子.框架.test import load_and_preprocess_data
from 多因子.框架.回测 import RLStrategy
from 多因子.框架.策略 import TradingEnv, A2CAgent


def backtest(test_data, model_path, scaler_path, features,acmodel, init_balance=1e6):

    # ================== 模型加载 ==================
    lstm_model, scaler = 模型训练.LSTM(test_data, features)

    # 加载训练好的权重
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型文件不存在: {model_path}")
    lstm_model.load_state_dict(torch.load(model_path))
    lstm_model.eval()

    # 加载归一化器
    import joblib
    if not os.path.exists(scaler_path):
        raise FileNotFoundError(f"归一化器文件不存在: {scaler_path}")
    scaler = joblib.load(scaler_path)

    def create_cerebro(env):
        cerebro = bt.Cerebro()

        # 重命名日期列为datetime
        df = test_data.reset_index().rename(columns={'trade_date': 'datetime'})
        df['datetime'] = pd.to_datetime(df['datetime'])  # 确保是datetime类型

        data = bt.feeds.PandasData(
            dataname=df[['datetime', 'open', 'high', 'low', 'close', 'volume']],
            datetime=0,  # 现在正确对应datetime列
            open=1,
            high=2,
            low=3,
            close=4,
            volume=5
        )
        cerebro.adddata(data)
        # AC智能体
        state_dim = 2  # 根据状态向量维度设置
        action_dim = 7
        agent = A2CAgent(state_dim, action_dim)
        agent.load(acmodel)  # 加载训练好的策略网络
        # 添加策略
        cerebro.addstrategy(RLStrategy, env=env, agent=agent, train_mode=False,verbose=True)

        # 资金配置
        cerebro.broker.setcash(init_balance)
        cerebro.broker.setcommission(commission=0.001)
        return cerebro

    # ================== 执行回测 ==================
    env = TradingEnv(
        data=test_data,
        lstm_model=lstm_model,
        scaler=scaler,
        features=features,
        init_balance=init_balance
    )
    print("================== 开始回测 ==================")
    cerebro = create_cerebro(env)
    results = cerebro.run()

    # ================== 结果分析 ==================
    strat = results[0]


    # 可视化结果
    plt.figure(figsize=(12, 6))
    cerebro.plot(style='candlestick', iplot=False)
    plt.title("Backtest Result")
    plt.savefig('./model/result/backtest_result.png')
    plt.close()

    return None


